package com.gmrz.webauthn.protocol.v1.schema;

import co.nstant.in.cbor.model.DataItem;
import com.gmrz.util.Convert;
import com.gmrz.webauthn.common.CBORDecodeUtil;
import com.gmrz.webauthn.common.CBOREncodeUtil;
import com.gmrz.webauthn.common.CBOREncodeValueWrapper;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnCborException;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnErrorCode;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnException;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.*;
import java.util.Map.Entry;


public class AuthenticatorData {
    private static final Logger LOG = LogManager.getLogger(AuthenticatorData.class);


    private static final int MINIMUM_AUTHNR_DATA_LENGTH = 37;


    private static final int MINIMUM_ATT_DATA_LENGTH = 22;


    private static final int ATTESTATION_DATA_INCLUDED = 64;


    private static final int EXTENSION_DATA_INCLUDED = 128;


    private byte[] rpID;


    private byte flags;


    private long signCount;


    private AttestedCredentialData attestation;


    private Map<String, Object> extensions;


    public byte[] getRpID() {
        return this.rpID;
    }

    public void setRpID(byte[] rpID) {
        this.rpID = rpID;
    }

    public byte getFlags() {
        return this.flags;
    }

    public void setFlags(byte flags) {
        this.flags = flags;
    }

    public long getSignCount() {
        return this.signCount;
    }

    public void setSignCount(long signCount) {
        this.signCount = signCount;
    }

    public AttestedCredentialData getAttestation() {
        return this.attestation;
    }

    public void setAttestation(AttestedCredentialData attestation) {
        this.attestation = attestation;
    }

    public Map<String, Object> getExtensions() {
        return this.extensions;
    }

    public void setExtensions(Map<String, Object> extensions) {
        this.extensions = extensions;
    }

    public void decode(byte[] rawData)
            throws WebAuthnCborException, WebAuthnException {

        ByteBuffer buffer = ByteBuffer.wrap(rawData).order(ByteOrder.BIG_ENDIAN);

        if (buffer.remaining() < 37) {
            String errorMsg = "Authenticator data length [" + rawData.length + "] is less than [" + 37 + "]";

            logAndThrowException(errorMsg);
        }


        this.rpID = new byte[32];

        buffer.get(this.rpID);

        this.flags = buffer.get();
        boolean tup = (this.flags & 0x1) > 0;

        boolean UV_included = (this.flags & 0x4) > 0;

        boolean AT_included = (this.flags & 0x40) > 0;
        boolean ED_included = (this.flags & 0x80) > 0;

        LOG.debug("TUP is " + (tup ? "" : "NOT ") + "set.");
        LOG.debug("User Verified : " + (UV_included ? "" : "NOT ") + "included.");
        LOG.debug("Attestation data is " + (AT_included ? "" : "NOT ") + "included.");
        LOG.debug("Extension data is " + (ED_included ? "" : "NOT ") + "included.");


        this.signCount = (buffer.getInt() & 0xFFFFFFFF);

        LOG.debug("Sign counter : {}", this.signCount);

        if (this.signCount < 0L) {
            String errorMsg = "Sign Counter must be greater than 0, value is : " + this.signCount;
            logAndThrowException(errorMsg);
        }
        List<DataItem> dataItems = null;

        if (AT_included) {
            if (buffer.remaining() < 22) {
                String errorMsg = "Attestation statement length [" + buffer.remaining() + "] is less than [" + 22 + "]";

                logAndThrowException(errorMsg);
            }

            this.attestation = new AttestedCredentialData();

            byte[] aaguidBytes = new byte[16];
            buffer.get(aaguidBytes);


            this.attestation.setAaguid(new AAGUID(aaguidBytes));


            short length = buffer.getShort();
            if (buffer.remaining() < length) {

                String errorMsg = "Credentail ID length [" + length + "] is greater than remaining [" + buffer.remaining() + "]";
                logAndThrowException(errorMsg);
            }
            byte[] credentialID = new byte[length];
            buffer.get(credentialID);

            this.attestation.setCredentialID(Convert.toBase64(credentialID));


            dataItems = getDataItems(buffer);

            co.nstant.in.cbor.model.Map cosePublicKEycborMap = (co.nstant.in.cbor.model.Map) dataItems.get(0);
            Map<String, Object> pubKeymap = CBORDecodeUtil.cborMapConversion(cosePublicKEycborMap);

            this.attestation.parseCosePublicKey(pubKeymap);
        }


        if (ED_included) {
            if (dataItems == null) {

                dataItems = getDataItems(buffer);
            }


            if ((dataItems != null) && (!dataItems.isEmpty())) {
                co.nstant.in.cbor.model.Map cborExtensionMap = (co.nstant.in.cbor.model.Map) dataItems.get(dataItems.size() - 1);

                Map<String, Object> extens = CBORDecodeUtil.cborMapConversion(cborExtensionMap);
                LOG.debug("Extension data is" + extens.toString());
            }
        } else if (buffer.hasRemaining()) {
            String errorMsg = "Extra bytes found during decoding the authenticator data.";
            logAndThrowException(errorMsg);
        }
    }

    private List<DataItem> getDataItems(ByteBuffer buffer)
            throws WebAuthnCborException, WebAuthnException {
        int remainingBufferLen = buffer.remaining();

        if (remainingBufferLen == 0) {
            String errorMsg = "buffer doesn't have enough remaining for public key or extenstion data";

            logAndThrowException(errorMsg);
        }
        byte[] reaminingData = new byte[remainingBufferLen];
        buffer.get(reaminingData);


        return CBORDecodeUtil.cborDecode(reaminingData);
    }

    public byte[] encode() throws WebAuthnException {
        LOG.debug("Encoding authenticator data bytes");

        if (this.rpID == null) {
            String errorMsg = "Cannot encode to bytes with RPID null";
            logAndThrowException(errorMsg);
        }
        if (this.signCount == 0L) {
            String errorMsg = "Cannot encode to bytes with signature counter null";
            logAndThrowException(errorMsg);
        }

        byte[] attestationData = null;
        if (this.attestation != null) {
            this.flags = ((byte) (this.flags | 0x40));

            attestationData = this.attestation.encodeToBytes();
        }

        byte[] extensionBytes = new byte[0];
        if (this.extensions != null) {
            this.flags = ((byte) (this.flags | 0x80));
            try {
                extensionBytes = CBOREncodeUtil.cborMapEncode(createExtensionMap(this.extensions));
            } catch (WebAuthnCborException e) {
                String errorMsg = "Could note CBOR encode extension bytes";
                throw new WebAuthnException(WebAuthnErrorCode.ERROR_CBOR_ENCODING_FAILED, errorMsg);
            }
        }


        int authDatalength = 37;
        int attDataLength = attestationData != null ? attestationData.length : 0;
        int extDatalength = extensionBytes != null ? extensionBytes.length : 0;

        byte[] authenticatorData = new byte[authDatalength + attDataLength + extDatalength];


        int intSignCount = (int) this.signCount;
        ByteBuffer.wrap(authenticatorData).put(this.rpID).put(this.flags).putInt(intSignCount);

        if (attDataLength != 0) {
            ByteBuffer.wrap(authenticatorData, authDatalength, attDataLength).put(attestationData);
        }
        if (extDatalength != 0) {
            ByteBuffer.wrap(authenticatorData, authDatalength + attDataLength, extDatalength).put(extensionBytes);
        }

        return authenticatorData;
    }


    private Map<CBOREncodeValueWrapper, CBOREncodeValueWrapper> createExtensionMap(Map<String, Object> extensionData)
            throws WebAuthnException {
        Map<CBOREncodeValueWrapper, CBOREncodeValueWrapper> cborExtensionMap = new HashMap<CBOREncodeValueWrapper, CBOREncodeValueWrapper>();

        for (Entry<String, Object> entry : extensionData.entrySet()) {
            Object extensionId = entry.getKey();
            Object extensionValue = entry.getValue();

            if (extensionValue.getClass() == String.class) {
                cborExtensionMap.put(new CBOREncodeValueWrapper(extensionId, String.class), new CBOREncodeValueWrapper(extensionValue, String.class));
            } else if (extensionValue.getClass() == Map.class) {
                Map<CBOREncodeValueWrapper, CBOREncodeValueWrapper> cborExtensionValueMap = new HashMap<CBOREncodeValueWrapper, CBOREncodeValueWrapper>();

                Map<Object, Object> valueMap = (Map) extensionValue;


                for (Entry<Object, Object> innerEntry : valueMap.entrySet()) {
                    Object key = innerEntry.getKey();
                    Object value = innerEntry.getValue();


                    if (value.getClass() == ArrayList.class) {
                        ArrayList<CBOREncodeValueWrapper> wrapperArray = new ArrayList<CBOREncodeValueWrapper>();

                        for (Object val : (ArrayList) value) {
                            wrapperArray.add(new CBOREncodeValueWrapper(val, val.getClass()));
                        }
                        cborExtensionValueMap.put(new CBOREncodeValueWrapper(key, String.class), new CBOREncodeValueWrapper(wrapperArray, ArrayList.class));
                    } else {
                        cborExtensionValueMap.put(new CBOREncodeValueWrapper(key, String.class), new CBOREncodeValueWrapper(value, value.getClass()));
                    }
                }

                cborExtensionMap.put(new CBOREncodeValueWrapper(extensionId, String.class), new CBOREncodeValueWrapper(cborExtensionValueMap, Map.class));
            } else {
                String errorMsg = "Invalid type of extension data";
                LOG.error(errorMsg);
                throw new WebAuthnException(WebAuthnErrorCode.STATUS_INVALID_EXTENSION, errorMsg);
            }
        }

        return cborExtensionMap;
    }

    private void logAndThrowException(String errorMsg) throws WebAuthnException {
        LOG.error(errorMsg);
        throw new WebAuthnException(WebAuthnErrorCode.STATUS_INVALID_AUTHENTICATOR_DATA, errorMsg);
    }


    public String toString() {
        return "AuthenticatorData [rpID=" + Arrays.toString(this.rpID) + ", tup=" + this.flags + ", signCount=" + this.signCount + ", attestation=" + this.attestation + ", extensions=" + this.extensions + "]";
    }
}
